#include "stdafx.h"
#include "Flocking.h"
using namespace std;

Flock::Flock( )
{
	this->alignmentStrength = 0;
	this->alignmentType = 0;
	this->boundaryType = 0;
	this->clusteringEnabled = false;
	this->cohesionStrength = 0;
	this->cohesionType = 0;
	this->gridEnabled = false;
	this->initRange = 0;
	this->lambda = 0;
	this->nBirds = 0;
	this->nClusters = 0;
	this->neighborNum = 0;
	this->neighborRange = 0;
	this->nGrids = 0;
	this->noiseType = 0;
	this->noiseLevel = 0;
	this->systemSize = 0;
	this->timeStep = 0;
	this->c_rb = 0;
	this->c_ra = 0;
	this->c_re = 0;
	this->c_rS = 0;
	this->c_aS = 0;
}

Flock::Flock( Configuration& config)
{
	if ( config.configurationEnabled == false)
	{
		cout<<"Flock: configuration not enabled."<<endl;
	}
	else
	{
		this->alignmentStrength = config.alignmentStrength;
		this->alignmentType = config.alignmentType;
		this->boundaryType = config.boundaryType;
		this->clusteringEnabled = config.clusteringEnabled;
		this->cohesionStrength = config.cohesionStrength;
		this->cohesionType = config.cohesionType;
		this->gridEnabled = config.gridEnabled;
		this->initRange = config.initRange;
		this->lambda = config.weightVelocity;
		this->nBirds = config.flockSize;
		this->nClusters = config.clusterNum;
		this->neighborNum = config.neighborNum;
		this->neighborRange = config.neighborRange;
		this->nGrids = config.systemSize / (2*config.neighborRange);
		this->noiseType = config.noiseType;
		this->noiseLevel = config.noiseStrength;
		this->systemSize = config.systemSize;
		this->timeStep = config.timeStep;
		this->c_rb = config.c_rb;
		this->c_ra = config.c_ra;
		this->c_re = config.c_re;
		this->c_rS = config.c_rS;
		this->c_aS = config.c_aS;
	}
}


Flock::~Flock()
{
	delete [] bird;   
	bird   = NULL;
	if ( gridEnabled == true )
	{
		delete [] grid;    
		grid   = NULL;
	}
	if ( clusteringEnabled == true)
	{
		delete [] center;  
		center = NULL;
		delete [] color;   
		color  = NULL;
	}
}

void Flock::addIntrinsicNoise( )
{
	double theta;
	double phi;
	Vector_3D X;
	std::tr1::mt19937 eng;
	std::tr1::normal_distribution<double> dist;
	std::tr1::uniform_real<double> unif;
	eng.seed((unsigned int)0);
	for (int id = 0; id < nBirds; id++)
	{
		switch ( this->noiseType )
		{
		case GAUSSIAN : { theta = noiseLevel*dist(eng); break; }
		case UNIFORM : { theta = pi*noiseLevel * unif(eng);break;}
		}
		phi = 2*pi*unif(eng);
		X = rotateTowardZaxis(bird[id].V, theta);
		X = rotateAroundOrigin(X, bird[id].V, phi);
		bird[id].V = X;
	}
}



void Flock::boundary( )
{
	if (this->boundaryType == OPEN ) return;
	if ( this->boundaryType  == PERIODIC )
	{
		for (int id = 0; id < nBirds; id++)
		{
			bird[id].X.x = bird[id].X.x < 0 ? bird[id].X.x + this->systemSize : bird[id].X.x;
			bird[id].X.x = bird[id].X.x > this->systemSize ? bird[id].X.x - this->systemSize : bird[id].X.x;
			bird[id].X.y = bird[id].X.y < 0 ? bird[id].X.y + this->systemSize : bird[id].X.y;
			bird[id].X.y = bird[id].X.y > this->systemSize ? bird[id].X.y - this->systemSize : bird[id].X.y;
			bird[id].X.z = bird[id].X.z < 0 ? bird[id].X.z + this->systemSize : bird[id].X.z;
			bird[id].X.z = bird[id].X.z > this->systemSize ? bird[id].X.z - this->systemSize : bird[id].X.z;
		}
		return ;
	}
	if ( this->boundaryType == REFLEXIVE )
	{
		for (int id = 0; id < nBirds; id++)
		{
			if ( bird[id].X.x < 0 )
			{
				bird[id].X.x = -2*bird[id].X.x;
				bird[id].V.x = -bird[id].V.x;
			}
			if( bird[id].X.x > this->systemSize)
			{
				bird[id].X.x = 2*this->systemSize - bird[id].X.x;
				bird[id].V.x = - bird[id].V.x;
			}
			if ( bird[id].X.y < 0 )
			{
				bird[id].X.y = -2*bird[id].X.y;
				bird[id].V.y = -bird[id].V.y;
			}
			if( bird[id].X.y > this->systemSize)
			{
				bird[id].X.y = 2*this->systemSize - bird[id].X.y;
				bird[id].V.y = - bird[id].V.y;
			}
			if ( bird[id].X.z < 0 )
			{
				bird[id].X.z = -2*bird[id].X.z;
				bird[id].V.z = -bird[id].V.z;
			}
			if( bird[id].X.z > this->systemSize)
			{
				bird[id].X.z = 2*this->systemSize - bird[id].X.z;
				bird[id].V.z = - bird[id].V.z;
			}
		}
		return ;
	}
}

inline double Flock::distance(Ball_3D* first, Ball_3D* second)
{
	return (first->X - second->X).norm();
}


Vector_3D* Flock::getAlignmentForce( )
{
	Vector_3D* force = new Vector_3D [nBirds];
	if ( this->alignmentType == TOPOLOGICAL )
	{
		for (int id = 0; id < nBirds; id++)
		{
			int* it = getNearestNeighbors_Full( id );
			for ( int i = 0; i < this->neighborNum; i++)
			{
				if ((*it) < this->neighborNum && (*it)>=0)
				{
				double s = bird[it[i]].V.x* bird[id].V.x + bird[it[i]].V.y* bird[id].V.y +bird[it[i]].V.y* bird[id].V.y ;
				force[id] = force[id] + bird[it[i]].V/this->neighborNum + bird[id].V/ this->neighborNum * bird[it[i]].V.norm()*bird[it[i]].V.norm();
				}
			}
			//force[id] = (force[id] / this->neighborNum - bird[id].V);
			delete [] it;
		}
	}

	//if ( this->alignmentType == METRIC )
	//{
	//	for(int id = 0; id < nBirds; id++)
	//	{
	//		set<int> neighbor;
	//		if ( this->gridEnabled == false) 
	//		{
	//			neighbor = getNeighborId_Full( id );
	//		}
	//		else
	//		{
	//			neighbor = getNeighborId( id );
	//		}
	//		for(set<int>::iterator it = neighbor.begin(); it != neighbor.end(); it++)
	//		{
	//			force[id] = force[id] + bird[*it].V;
	//		}
	//		force[id] = (force[id] / neighbor.size() - bird[id].V);
	//	}
	//}
	for (int id = 0; id < this->nBirds; id++)
			force[id] = force[id] * this->alignmentStrength;
	return force;
}

Vector_3D* Flock::getCohesionForce()
{
	Vector_3D* force = new Vector_3D [nBirds];
	if ( this->cohesionType == LOCAL )
	{
		for(int id = 0; id < this->nBirds; id++)
		{
			set<int> neighbor;
			if ( this->gridEnabled == false) 
			{
				neighbor = getNeighborId_Full( id );
			}
			else
			{
				neighbor = getNeighborId( id );
			}
			for(set<int>::iterator it = neighbor.begin(); it != neighbor.end(); it++)
			{
				force[id] = force[id] + pairCohensionForce(bird[id], bird[*it]);
			}
		}
	}
	if ( this->cohesionType == GLOBAL )
	{
		//for(int id = 0; id < this->nBirds; id++)
		//      for (int j = 0; j < this->nBirds; j++)
		//              if ( j != id) force[id] = force[id] + pairCohensionForce(bird[id], bird[j]);
		for (int id = 0; id < nBirds; id++)
		{
			int* it = getNearestNeighbors_Full( id );
			for ( int i = 0; i < this->neighborNum; i++)
			{
				force[id] = force[id] + pairCohensionForce(bird[id], bird[it[i]]);
			}
			delete [] it;
		}

	}
	for (int i = 0; i < this->nBirds; i++)
		force[i] = force[i] * cohesionStrength;
	return force;
}

Vector_3D* Flock::getExtrinsicNoise()
{
	Vector_3D* force = new Vector_3D [nBirds];
	double theta;
	double phi;
	std::tr1::mt19937 eng;
	std::tr1::normal_distribution<double> dist;
	std::tr1::uniform_real<double> unif;
	eng.seed((unsigned int)0);
	Vector_3D X;
	for (int id = 0; id < nBirds; id++)
	{
		switch ( noiseType )
		{
		case GAUSSIAN : theta = noiseLevel*dist(eng); break;
		case UNIFORM : theta = pi*noiseLevel * unif(eng);break;
		}
		phi = 2*pi*unif(eng);
		X = rotateTowardZaxis(bird[id].V, theta);
		force[id] = rotateAroundOrigin(X, bird[id].V, phi);
	}
	return force;
}


int* Flock::getNearestNeighbors( int id)
{
	if (this->gridEnabled == false )
	{
		cout<<"Flock: grid not enabled while use getNearestNeighbors"<<endl;
	}
	double  gridsize = (double)this->systemSize / this->nGrids;
	int j = floor(bird[id].X.x / gridsize);
	int k = floor(bird[id].X.y / gridsize);
	int l = floor(bird[id].X.z / gridsize);
	int x1[3] = {(j-1) < 0 ? (j-1+nGrids) % nGrids : (j-1) % nGrids, j, (j+1) % nGrids};
	int x2[3] = {(k-1) < 0 ? (k-1+nGrids) % nGrids : (k-1) % nGrids, k, (k+1) % nGrids};
	int x3[3] = {(l-1) < 0 ? (l-1+nGrids) % nGrids : (l-1) % nGrids, l, (l+1) % nGrids};
	int count = 0;
	for (j = 0;j < 3;j++) {
		for(k = 0;k < 3;k++) {
			for(l = 0;l < 3;l++)
			{
				int offset = x3[l]*nGrids*nGrids + x2[k]*nGrids + x1[j];
				count+=grid[offset].size();
			}
		}
	}

	if ( count < this->neighborNum )
	{
		cout<<"Flock : find neighbors globally"<<endl;
		return getNearestNeighbors_Full( id );
	}
	double* dist = new double [count];
	int* neighborId = new int [count];
	int* nearestNeighborId = new int [this->neighborNum];
	int cnt = 0;
	for (j = 0;j < 3;j++) {
		for(k = 0;k < 3;k++) {
			for(l = 0;l < 3;l++)
			{
				int offset = x3[l]*nGrids*nGrids + x2[k]*nGrids + x1[j];
				for(set<int>::iterator it = grid[offset].begin(); it != grid[offset].end(); it++)
				{
					neighborId[cnt] = *it;
					Vector_3D v(bird[id].X - bird[*it].X);
					dist[cnt] = v.x*v.x + v.y*v.y + v.z*v.z;
					cnt++;
				}
			}
		}
	}
	double distK = findKthMin(dist, 0, count-1, this->neighborNum);
	cnt = 0;
	for (int i = 0; i < count; i++)
	{
		if ( dist[i] <= distK + 0.00001)
		{
			nearestNeighborId[cnt] = neighborId[i];
			cnt++;
		}
	}
	delete [] dist;
	delete [] neighborId;
	return nearestNeighborId;
}


int* Flock::getNearestNeighbors_Full( int id)
{
	int* it = new int [this->neighborNum];
	double* dist = new double [nBirds];
	double* dist2 = new double [nBirds];
	for (int i = 0; i < this->nBirds; i++)
	{
		Vector_3D v(bird[i].X - bird[id].X);
		dist[i] = v.x*v.x+v.y*v.y+v.z*v.z;
		dist2[i] = dist[i];
	}
	double distK = findKthMin(dist, 0, this->nBirds-1, this->neighborNum);
	int cnt = 0;
	for (int i = 0; i < nBirds; i++)
	{
		if ( dist2[i] <= distK + 0.00001)
		{
			it[cnt] = i;
			cnt++;
		}
		if (cnt == this->neighborNum) break;
	}
	if ( cnt < this->neighborNum) 
	{
		cout<<" Flock : can not find "<<this->neighborNum;
		cout<<" in getting nearest neighbors"<<endl;
	}
	delete [] dist;
	delete [] dist2;
	return it;
}


set<int> Flock::getNeighborId( int id)
{
	if (this->gridEnabled == false )
	{
		cout<<"Flock: grid not enabled in getNeighborId"<<endl;
	}
	double gridsize = this->systemSize / this->nGrids;
	set<int> neighbor;
	int j = floor(bird[id].X.x / gridsize);
	int k = floor(bird[id].X.y / gridsize);
	int l = floor(bird[id].X.z / gridsize);
	int x1[3] = {(j-1) < 0 ? (j-1+nGrids) % nGrids : (j-1) % nGrids, j, (j+1) % nGrids};
	int x2[3] = {(k-1) < 0 ? (k-1+nGrids) % nGrids : (k-1) % nGrids, k, (k+1) % nGrids};
	int x3[3] = {(l-1) < 0 ? (l-1+nGrids) % nGrids : (l-1) % nGrids, l, (l+1) % nGrids};
	for (j = 0;j < 3;j++) {
		for(k = 0;k < 3;k++) {
			for(l = 0;l < 3;l++)
			{
				int offset = x3[l]*nGrids*nGrids + x2[k]*nGrids + x1[j];
				for(set<int>::iterator it = grid[offset].begin(); it != grid[offset].end(); it++)
				{
					int tid = *it;
					Vector_3D v(bird[tid].X - bird[id].X);
					double dist = v.x*v.x+v.y*v.y+v.z*v.z;
					if (dist < neighborRange) neighbor.insert(tid);
				}
			}
		}
	}
	if ( neighbor.size() < 1)
	{
		cout<<"Flock : error, return empty neighbor"<<endl;
	}
	return neighbor;
}


set<int> Flock::getNeighborId_Full(int id)
{
	double dist = 0;
	set<int> neighbor;
	for (int i = 0; i < nBirds; i++)
	{
		Vector_3D v(bird[i].X - bird[id].X);
		dist = v.x*v.x + v.y*v.y + v.z*v.z;
		if ( dist < neighborRange )
			neighbor.insert(i);
	}
	return neighbor;
}

inline Vector_3D Flock::pairCohensionForce(Ball_3D& first, Ball_3D& second)
{
	Vector_3D X = second.X - first.X;
	double r = X.norm();
	double strength = 0;
	if ( &first == &second) return Vector_3D(0,0,0);
	//if ( r < c_rb ) strength = c_aS - c_ra / (c_rb * c_rb);
	//if ( r > c_rb ) strength = c_aS - c_ra / (r * r);
	      if ( r > c_re ) strength = 0;
	      if ( r > c_ra && r < c_re) strength = c_aS;
	      if ( r > c_rb && r < c_ra) strength = (c_aS + c_rS)/(c_ra - c_rb)*(r - c_rb) - c_rS;
	      if ( r < c_rb) strength = -c_rS;
	X = X/r;
	return X*strength;
}


void Flock::setConfiguration( Configuration & config)
{
	if ( config.configurationEnabled == false)
	{
		cout<<"Flock: configuration not enabled."<<endl;
	}
	else
	{
		this->alignmentStrength = config.alignmentStrength;
		this->alignmentType = config.alignmentType;
		this->boundaryType = config.boundaryType;
		this->clusteringEnabled = config.clusteringEnabled;
		this->cohesionStrength = config.cohesionStrength;
		this->cohesionType = config.cohesionType;
		this->gridEnabled = config.gridEnabled;
		this->initRange = config.initRange;
		this->lambda = config.weightVelocity;
		this->nBirds = config.flockSize;
		this->nClusters = config.clusterNum;
		this->neighborNum = config.neighborNum;
		this->neighborRange = config.neighborRange;
		this->nGrids = config.systemSize / (2*config.neighborRange);
		this->noiseType = config.noiseType;
		this->noiseLevel = config.noiseStrength;
		this->systemSize = config.systemSize;
		this->timeStep = config.timeStep;
		this->c_rb = config.c_rb;
		this->c_ra = config.c_ra;
		this->c_re = config.c_re;
		this->c_rS = config.c_rS;
		this->c_aS = config.c_aS;
	}
}

void Flock::updatePosition()
{
	for (int i = 0; i < nBirds; i++)
	{
		bird[i].X += bird[i].V * this->timeStep;
	}
	boundary();
}

bool Flock::init()
{
	if ( !initPV2() ) return false;
	if ( this->gridEnabled == true )                        
	{
		if( ( grid = new set<int> [nGrids * nGrids * nGrids]) == NULL )
			return false;
	}
	if ( this->clusteringEnabled == true)
		if ( !initCluster() ) return false;
	return true;
}

bool Flock::initCluster()
{
	if ( (color = new unsigned int [nBirds]) == NULL) return false;
	if ( (center = new Ball_3D [this->nClusters]) == NULL) return false;
	memset(color, 0, sizeof(unsigned int)*nBirds);
	std::tr1::mt19937 eng;
	uniform_int<int> unif(0, nBirds-1);
	eng.seed((unsigned int) 0);
	int cnt = 0;
	set<int> recId;
	set<int>::iterator it;
	if ( nBirds < nClusters ) return false;
	while( cnt < nClusters)
	{
		int id = unif(eng);
		if ((it=recId.find(id)) == recId.end() )
		{
			recId.insert(id);
			center[cnt] = bird[id];
			color[id] = cnt;
			cnt++;
		}
	}
	return true;
}

bool Flock::initPV()
{
	if ( (bird = new Ball_3D [nBirds]) == NULL ) return false;
	std::tr1::mt19937 eng;
	std::tr1::uniform_real<double> unif;
	eng.seed((unsigned int)0);
	double a,b,c,r;
	if ( this->initRange > this->systemSize)
	{
		cout<<"Flock: error while initializing position."<<endl;
		return false;
	}
	for (int i = 0; i < this->nBirds; i++)
	{
		a = this->initRange * (unif(eng) - 0.5) + this->systemSize * 0.5;
		b = this->initRange * (unif(eng) - 0.5) + this->systemSize * 0.5;
		c = this->initRange * (unif(eng) - 0.5) + this->systemSize * 0.5;
		bird[i].setPosition(a,b,c);
		a = unif(eng) - 0.5;
		b = unif(eng) - 0.5;
		c = unif(eng) - 0.5;
		r = sqrt(a*a + b*b + c*c);
		bird[i].setVelocity(a/r,b/r,c/r);
	}
	return true;
}

bool Flock::initPV2()
{
	if ( (bird = new Ball_3D [nBirds]) == NULL ) return false;
	std::tr1::mt19937 eng;
	std::tr1::uniform_real<double> unif;
	eng.seed((unsigned int)0);
	double a,b,c,r;
	if ( this->initRange > this->systemSize)
	{
		cout<<"Flock: error while initializing position."<<endl;
		return false;
	}
	for (int i = 0; i < this->nBirds; i++)
	{
		a = this->initRange * (unif(eng) - 0.5) + this->systemSize * 0.5;
		b = this->initRange * (unif(eng) - 0.5) + this->systemSize * 0.5;
		c = this->initRange * (unif(eng) - 0.5) + this->systemSize * 0.5;
		bird[i].setPosition(a,b,c);
		bird[i].setVelocity(1,0,0);
	}
	return true;
}

void Flock::kmeans()
{
	bool finished = false;
	while ( !finished )
	{
		unsigned int  variation = 0;
		for (int i = 0 ; i < nBirds; i++)
		{
			int k = nearestCenter(i);
			variation += abs((int)(color[i] - k));
			color[i] = k;
		}
		if ( variation == 0) finished = true;

		//update center position
		for(int i = 0; i < nClusters; i++)
			center[i] = Ball_3D(0,0,0,0,0,0,0,0);
		int* nbirdsInCluster = new int [nClusters];
		memset(nbirdsInCluster, 0, sizeof(int)*nClusters);
		for (int i = 0; i < nBirds; i++)
		{
			center[color[i]].X = center[color[i]].X + bird[i].X;
			center[color[i]].V = center[color[i]].V + bird[i].V;
			nbirdsInCluster[color[i]]++;
		}
		for (int i = 0; i < nClusters; i++)
		{
			center[i].X = center[i].X / nbirdsInCluster[i];
			center[i].V = center[i].V / nbirdsInCluster[i];
		}
		delete[] nbirdsInCluster;
	}
}


void Flock::mapToGrid()
{
	if ( this->gridEnabled == false )
	{
		cout<<"Flock: grid not enabled in mapToGrid"<<endl;
		return;
	}
	double gridsize = (double)systemSize / nGrids;
	for (int i = 0; i < nGrids*nGrids*nGrids; i++)
		grid[i].clear();
	for (int i = 0; i < nBirds; i++)
	{
		int j = floor(bird[i].X.x / gridsize);
		int k = floor(bird[i].X.y / gridsize);
		int l = floor(bird[i].X.z / gridsize);
		grid[l*nGrids*nGrids + k*nGrids + j].insert(i);
	}
}


void Flock::move()
{
	this->updateVelocity();
	this->updatePosition();
	if ( this->clusteringEnabled == true ) kmeans();
}


int Flock::nearestCenter( int id )
{
	int k = 0;
	double min = INF;
	for (int i = 0; i < nClusters; i++)
	{
		double dist = quasiDistance(&center[i], &bird[id]);
		if ( dist < min )
		{
			k = i;
			min = dist;
		}
	}
	return k;
}

void Flock::updateVelocity()
{
	if ( this->gridEnabled == true )
		this->mapToGrid();
	Vector_3D* force = new Vector_3D [nBirds];
	Vector_3D* aF = new Vector_3D [nBirds];
	Vector_3D* cF = new Vector_3D [nBirds];
	if ( abs(this->alignmentStrength) > 1e-8)  aF = this->getAlignmentForce();
	if ( abs(this->cohesionStrength)  > 1e-8 )  cF = this->getCohesionForce();
	for (int i = 0; i < nBirds; i++)
	{
		force[i] = aF[i] + cF[i];
		//cout<<aF[i].norm()<<" "<<cF[i].norm()<<endl;
		bird[i].V = bird[i].V + force[i]*this->timeStep;
		bird[i].V = bird[i].V / bird[i].V.norm();
	}
	//system("pause");
	delete [] aF;
	delete [] cF;
	delete [] force;
	this->addIntrinsicNoise();
}


inline double Flock::quasiDistance( Ball_3D* first, Ball_3D* second)
{
	return  (first->X - second->X).norm() + (first->V - second->V).norm() * lambda;
}


bool Flock::specialInit()
{
	bird = new Ball_3D [nBirds];
	grid = new set<int> [nGrids * nGrids * nGrids];
	std::tr1::mt19937 eng;
	std::tr1::uniform_real<double> unif;
	eng.seed((unsigned int)time(0));
	double a,b,c;
	for (int i = 0; i < nBirds / 2; i++)
	{
		a = this->systemSize/2 * unif(eng);
		b = this->systemSize * unif(eng);
		c = this->systemSize * unif(eng);
		bird[i].setPosition(a,b,c);
		a = unif(eng);
		b = 0;
		c = 0;
		bird[i].setVelocity(a,b,c);
	}
	for (int i = nBirds/2 + 1; i < nBirds; i++)
	{
		a = this->systemSize/2 * unif(eng) + this->systemSize/2;
		b = this->systemSize * unif(eng);
		c = this->systemSize * unif(eng);
		bird[i].setPosition(a,b,c);
		a = -unif(eng);
		b = 0;
		c = 0;
		bird[i].setVelocity(a,b,c);
	}
	return true;
}

void Flock::specialUpdateSpeed()
{
	for (int i = 0; i < nBirds; i++)
		bird[i].V = bird[i].V / bird[i].V.norm();
}
